import os
import argparse
import math

import numpy as np
from matplotlib import pyplot as plt

import Helper

parser = argparse.ArgumentParser()
parser.add_argument('--folder', help='folder', type=str, default='build')
parser.add_argument('-d', '--datetime', help='date time', type=str, default='')
parser.add_argument('-i', '--init', help='init step', type=int, default='0')
parser.add_argument('-f', '--final', help='final step', type=int, default='-1')

args = parser.parse_args()
folder = args.folder
fname = args.datetime
start_idx = args.init
final_idx = args.final

if os.path.exists(os.path.join(folder, f'{fname}.npy')):
    data = np.load(os.path.join(folder, f'{fname}.npy'))[start_idx:final_idx, :]
else:
    data = np.loadtxt(os.path.join(folder, f'{fname}.csv'))
    np.save(os.path.join(folder, f'{fname}.npy'), data)
    data = data[start_idx:final_idx, :]
ob_raw = data[:, :62]
act_raw = data[:, -12:]
dt = 0.01
t = np.arange(data.shape[0]) * dt
kp = 40.0
kd = 1.0
max_torque = 23.7
knee_gear_ratio = 1.0

action_mean = np.array([0.0, 0.8, -1.75,
                        0.0, 0.8, -1.75,
                        0.0, 0.8, -1.75,
                        0.0, 0.8, -1.75])
action_std = np.array([0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9])

if not os.path.exists(os.path.join(folder, 'figure')):
    os.mkdir(os.path.join(folder, 'figure'))
figure_folder = os.path.join(folder, 'figure', fname)
try:
    os.mkdir(figure_folder)
except FileExistsError:
    print(f'folder: {figure_folder} already exists')

# phase
phase = np.arctan2(ob_raw[:, 3], ob_raw[:, 4])
plt.figure(figsize=[4.0, 3.0], dpi=600)
plt.xlabel('time (s)')
plt.ylabel('phase (rad)')
# plt.twinx()
plt.plot(t, phase, 'b')
plt.savefig(os.path.join(figure_folder, 'phase.png'))
plt.close()

# output
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
ylim = [[action_mean[i] - action_std[i], action_mean[i] + action_std[i]] for i in range(3)]
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = 'p (rad)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    kwargs['ylim'] = ylim[i % 3]
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t, act_raw[:, i])
plt.savefig(os.path.join(figure_folder, 'jdes.png'))
plt.close()

# joint pos
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
ylim = [[action_mean[i] - action_std[i], action_mean[i] + action_std[i]] for i in range(3)]
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = 'p (rad)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    kwargs['ylim'] = ylim[i % 3]
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t, ob_raw[:, i+5])
plt.savefig(os.path.join(figure_folder, 'jpos.png'))
plt.close()

# joint pos
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
ylim = [[action_mean[i] - action_std[i], action_mean[i] + action_std[i]] for i in range(3)]
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = 'p (rad)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    kwargs['ylim'] = ylim[i % 3]
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t, ob_raw[:, i+5], 'r')
    plt.plot(t, act_raw[:, i], 'g', alpha=0.5)
plt.savefig(os.path.join(figure_folder, 'jpos_comp.png'))
plt.close()

plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
ylim = []
for i in range(3):
    ob_hi = (ob_raw[:, [5 + j * 3 + i for j in range(4)]]).max()
    ob_lo = (ob_raw[:, [5 + j * 3 + i for j in range(4)]]).min()
    act_hi = (act_raw[:, [j * 3 + i for j in range(4)]]).max()
    act_lo = (act_raw[:, [j * 3 + i for j in range(4)]]).min()
    hi = max(ob_hi, act_hi)
    lo = min(ob_lo, act_lo)
    ylim.append([lo - (hi - lo) * 0.08, hi + (hi - lo) * 0.08])
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = 'p (rad)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    kwargs['ylim'] = ylim[i % 3]
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t, ob_raw[:, i+5], 'r')
    plt.plot(t, act_raw[:, i], 'g', alpha=0.5)
plt.savefig(os.path.join(figure_folder, 'jpos_comp2.png'))
plt.close()

# joint vel
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
ylim = [[-5, 5], [-10, 10], [-10, 10]]
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = 'v (rad/s)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    kwargs['ylim'] = ylim[i % 3]
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t, ob_raw[:, i+17])
    plt.plot(t[:-1], np.diff(ob_raw[:, i+5])/dt, '--', alpha=0.5)
plt.savefig(os.path.join(figure_folder, 'jvel.png'))
plt.close()

# joint acc
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        kwargs['yticks'] = []
    else:
        kwargs['ylabel'] = r'a (rad/$s^2$)'
    if subplotpos < 9:
        kwargs['xticks'] = []
    else:
        kwargs['xlabel'] = 't (s)'
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(t[:-1], np.abs(np.diff(ob_raw[:, i+17]))/dt)
    # plt.plot(t[:-1], np.diff(ob_raw[:, i+5])/dt, '--', alpha=0.5)
plt.savefig(os.path.join(figure_folder, 'jacc.png'))
plt.close()

# joint acc hist
plt.figure(figsize=[8.0, 6.0], dpi=600)
for i in range(12):
    plt.subplot(3, 4, Helper.remap_index(i)+1)
    plt.hist(np.abs(np.diff(ob_raw[:, i+17]))/dt)
    # plt.plot(t[:-1], np.diff(ob_raw[:, i+5])/dt, '--', alpha=0.5)
plt.savefig(os.path.join(figure_folder, 'jacc_hist.png'))
plt.close()

# body pose
fig = plt.figure(figsize=[4.0, 3.0], dpi=600)
ax = fig.add_axes([0.2, 0.2, 0.75, 0.7])
ax.plot(t, ob_raw[:, -3], 'r--', alpha=0.2)
ax.plot(t, ob_raw[:, 29], 'r')
ax.legend(['command', 'state'], loc='upper right')
# ax.set_ylim([-1.0, 1.0])
ax.set_ylabel(r'$z_x^b$')
ax.set_xlabel('t (s)')
ax.set_title('pitch')
plt.savefig(os.path.join(figure_folder, 'pitch.png'))
plt.close()

fig = plt.figure(figsize=[4.0, 3.0], dpi=600)
ax = fig.add_axes([0.2, 0.2, 0.75, 0.7])
ax.plot(t, ob_raw[:, -2], 'r--', alpha=0.2)
ax.plot(t, ob_raw[:, 30], 'r')
ax.legend(['command', 'state'], loc='upper right')
# ax.set_ylim([-1.0, 1.0])
ax.set_title('roll')
ax.set_ylabel(r'$z_y^b$')
ax.set_xlabel('t (s)')
plt.savefig(os.path.join(figure_folder, 'roll.png'))
plt.close()

# cmd
plt.figure(figsize=[8.0, 6.0], dpi=600)
plt.plot(t, ob_raw[:, 0], 'r')
plt.plot(t, ob_raw[:, 1], 'g')
plt.ylabel('v (m/s)')
plt.xlabel('time (s)')
plt.legend([r'$v_x$', r'$v_y$'], loc='upper left')
plt.twinx()
plt.plot(t, ob_raw[:, 2], 'b')
plt.ylabel(r'$\omega$ (rad/s)')
plt.legend([r'$\omega$'], loc='upper right')
plt.savefig(os.path.join(figure_folder, 'cmd.png'))

for i in range(4):
    plt.figure(figsize=[8.0, 6.0], dpi=600)
    plt.plot(t, ob_raw[:, 35+3*i], 'r')
    plt.plot(t, ob_raw[:, 35+3*i+1], 'g')
    if ob_raw.shape[1] == 62:
        plt.plot(t, ob_raw[:, 47+3*i], 'r--')
        plt.plot(t, ob_raw[:, 47+3*i+1], 'g--')
    ylo, yhi = plt.ylim()
    plt.ylabel('horizontal pos (m)')
    plt.xlabel('time (s)')
    if i == 0 or i == 3:
        plt.fill_between(t, -1.2, 1.2, where=phase > 0, color='green', alpha=0.5)
    else:
        plt.fill_between(t, -1.2, 1.2, where=phase < 0, color='green', alpha=0.5)
    plt.legend([r'$x^{dest}_{stance}$', r'$y^{dest}_{stance}$', r'$x^{dest}_{swing}$', r'$y^{dest}_{swing}$', r'$x^{pos}$', r'$y^{pos}$'], loc='upper left', ncol=3)
    plt.ylim([ylo, yhi])
    plt.savefig(os.path.join(figure_folder, f'foot_des_{Helper.get_joint_name(i*3)[0]}_horizontal.png'))

    plt.figure(figsize=[8.0, 6.0], dpi=600)
    plt.plot(t, ob_raw[:, 35+3*i+2], 'b')
    if ob_raw.shape[1] == 62:
        plt.plot(t, ob_raw[:, 47+3*i+2], 'b--')
    plt.ylabel('vertical pos (m)')
    plt.xlabel('time (s)')
    ylo, yhi = plt.ylim()
    if i == 0 or i == 3:
        plt.fill_between(t, -1.2, 1.2, where=phase > 0, color='green', alpha=0.5)
    else:
        plt.fill_between(t, -1.2, 1.2, where=phase < 0, color='green', alpha=0.5)
    plt.legend([r'$z^{dest}_{stance}$', r'$z^{dest}_{swing}$', r'$z^{pos}$'], loc='upper right')
    plt.ylim([ylo, yhi])
    plt.savefig(os.path.join(figure_folder, f'foot_des_{Helper.get_joint_name(i*3)[0]}_vertical.png'))

# tau est
tau_est = kp * (act_raw - ob_raw[:, 5:17]) - kd * ob_raw[:, 17:29]
plt.figure(figsize=[8.0, 6.0], dpi=600)
plt.xlabel('time (s)')
plt.ylabel('torque (Nm)')
# plt.twinx()
for i in range(4):
    plt.plot(t, tau_est[:, 3*i], 'r')
    plt.plot(t, tau_est[:, 3*i+1], 'g')
    plt.plot(t, tau_est[:, 3*i+2], 'b--')
plt.plot([t[0], t[-1]], [-max_torque, -max_torque], 'r', alpha=0.4)
plt.plot([t[0], t[-1]], [max_torque, max_torque], 'r', alpha=0.4)
plt.plot([t[0], t[-1]], [-max_torque*knee_gear_ratio, -max_torque*knee_gear_ratio], 'r', alpha=0.4)
plt.plot([t[0], t[-1]], [max_torque*knee_gear_ratio, max_torque*knee_gear_ratio], 'r', alpha=0.4)
# plt.legend(['phase'], loc='upper right')
# plt.ylabel('phase (rad)')
plt.legend(['abductor', 'hip', 'knee'], loc='lower right')
plt.savefig(os.path.join(figure_folder, 'torque.png'))
plt.close()

# wc
plt.figure(figsize=[8.0, 6.0], dpi=600, tight_layout={'pad': 1.5, 'h_pad': 0.8, 'w_pad': 0.6})
tau_est[:, 2] /= knee_gear_ratio
tau_est[:, 5] /= knee_gear_ratio
tau_est[:, 8] /= knee_gear_ratio
tau_est[:, 11] /= knee_gear_ratio
v = ob_raw[:, 17:29].copy()
v[:, 2] *= knee_gear_ratio
v[:, 5] *= knee_gear_ratio
v[:, 8] *= knee_gear_ratio
v[:, 11] *= knee_gear_ratio
for i in range(12):
    subplotpos = Helper.remap_index(i)+1
    kwargs = {}
    if subplotpos % 4 != 1:
        # kwargs['yticks'] = []
        pass
    else:
        kwargs['ylabel'] = r'$\tau$ ($N\cdot m$)'
    if subplotpos < 9:
        # kwargs['xticks'] = []
        pass
    else:
        kwargs['xlabel'] = 'v (rad/s)'
    plt.subplot(3, 4, subplotpos, **kwargs)
    plt.plot(v[:, i], tau_est[:, i])
    plt.xlim([-18, 18])
    plt.ylim([-max_torque, max_torque])
    plt.grid(True)
plt.savefig(os.path.join(figure_folder, 'wc.png'))
plt.close()